import argparse

import torch
from transformers import pipeline, AutoTokenizer

parser = argparse.ArgumentParser(add_help=True, description='lijing')
parser.add_argument('--model_name_or_path', default="hfl/chinese-roberta-wwm-ext", type=str, help='lujing')
parser.add_argument('--reward_model_name_or_path', default="hfl/chinese-roberta-wwm-ext", type=str, help='lujing')
parser.add_argument('--batch_size', default=4, type=int, help='lujing')
parser.add_argument('--mini_batch_size', default=1, type=int, help='lujing')
parser.add_argument('--num_epochs', default=100, type=int, help='lujing')
parser.add_argument('--save_interval', default=100, type=int, help='lujing')
parser.add_argument('--save_dir', default="./save_model/", type=str, help='lujing')
parser.add_argument('--local_rank', default=0, type=int, help='lujing')
parser.add_argument('--text_column', default="Tweet text", type=str, help='lujing')
parser.add_argument('--label_column', default="text_label", type=str, help='lujing')
parser.add_argument('--seed', default=0, type=int, help='lujing')
parser.add_argument('--steps', default=20000, type=int, help='lujing')
parser.add_argument('--learning_rate', default=1.41e-5, type=float, help='lujing')
parser.add_argument('--log_with', default=None, type=str, help='lujing')
parser.add_argument('--gradient_accumulation_steps', default=4, type=int, help='lujing')
parser.add_argument('--optimize_cuda_cache', default=True, type=bool, help='lujing')
parser.add_argument('--early_stopping', default=False, type=bool, help='lujing')
parser.add_argument('--output_max_length', default=128, type=int, help='lujing')
parser.add_argument('--target_kl', default=0.1, type=float, help='lujing')
parser.add_argument('--ppo_epochs', default=4, type=int, help='lujing')
parser.add_argument('--max_length', default=512, type=int, help='lujing')
parser.add_argument('--lora_alpha', default=16, type=int, help='lujing')
parser.add_argument('--lora_dropout', default=0.1, type=int, help='lujing')
parser.add_argument('--init_kl_coef', default=0.2, type=float, help='lujing')
parser.add_argument('--adap_kl_ctrl', default=True, type=bool, help='lujing')
parser.add_argument('--reward_baseline', default=0.0, type=float, help='lujing')
parser.add_argument('--save_freq', default=None, type=int, help='lujing')
args = parser.parse_args()


args.reward_model_name_or_path = "D:\\model_path\\roberta-base"

tokenizer = AutoTokenizer.from_pretrained(args.reward_model_name_or_path)
# Need to do this for gpt2, because it doesn't have an official pad token.
tokenizer.pad_token = tokenizer.eos_token

reward_kwargs = {
    "return_all_scores": True,
    "function_to_apply": "none",
    "batch_size": 16,
    "truncation": True,
    "max_length":args.max_length,
    # "return_tensors":"pt"
    # "pad_token":tokenizer.eos_token
}

sentiment_pipe = pipeline(
    task="sentiment-analysis",
    model=args.reward_model_name_or_path,

    tokenizer=tokenizer,
    return_token_type_ids=False,
)
texts=[]
texts.append("Question: When ever my home page in **`angular 2(ionic 2)`** app is loaded I want call service/function. How to achieve this?For the first time when the app is loaded(home page is loaded) I can write this in the **`constructor`**,"
             " but when user start using the app and **`push`** new pages into **`nav controller`** and **`pop`**"
             " them to come back to **`home page`**, the **`constructor`** wont get called again.### I am stuck here."
             "I would like to know what is the best way to achieve this feature?I am n00b in"
             " `angular 2` and `ionic 2`Answer: You can use the `ngOnInit()` method, which will be called on every view"
             " change.\begin{code}import { Component, OnInit } from '@angular/core';@Component({selector: 'my-app',"
             " but when user start using the app and **`push`** new pages into **`nav controller`** and **`pop`**"
             " them to come back to **`home page`**, the **`constructor`** wont get called again.### I am stuck here."
             "I would like to know what is the best way to achieve this feature?I am n00b in"
             " `angular 2` and `ionic 2`Answer: You can use the `ngOnInit()` method, which will be called on every view"
             " change.\begin{code}import { Component, OnInit } from '@angular/core';@Component({selector: 'my-app',"
             " but when user start using the app and **`push`** new pages into **`nav controller`** and **`pop`**"
             " them to come back to **`home page`**, the **`constructor`** wont get called again.### I am stuck here."
             "I would like to know what is the best way to achieve this feature?I am n00b in"
             " `angular 2` and `ionic 2`Answer: You can use the `ngOnInit()` method, which will be called on every view"
             " change.\begin{code}import { Component, OnInit } from '@angular/core';@Component({selector: 'my-app',"
             " but when user start using the app and **`push`** new pages into **`nav controller`** and **`pop`**"
             " them to come back to **`home page`**, the **`constructor`** wont get called again.### I am stuck here."
             "I would like to know what is the best way to achieve this feature?I am n00b in"
             " `angular 2` and `ionic 2`Answer: You can use the `ngOnInit()` method, which will be called on every view"
             " change.\begin{code}import { Component, OnInit } from '@angular/core';@Component({selector: 'my-app',"
             )

pipe_outputs = sentiment_pipe(texts, **reward_kwargs)
print(pipe_outputs)